{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a682ea0b",
   "metadata": {},
   "source": [
    "# BentoML TensorFlow2 MNIST Tutorial\n",
    "\n",
    "Link to source code: https://github.com/bentoml/BentoML/tree/main/examples/tensorflow2_mnist/\n",
    "\n",
    "The code is based on the TensorFlow2 example code here: https://www.tensorflow.org/tutorials/quickstart/advanced\n",
    "\n",
    "Install required dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad00863",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84988464",
   "metadata": {
    "raw_mimetype": "text/markdown"
   },
   "source": [
    "If you are running MacOS use the following pip command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31a9cdfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements-macos.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45393b74",
   "metadata": {},
   "source": [
    "## Define the model\n",
    "\n",
    "First let's initiate the dataset we'll be using and then create a Model which we will use to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef7a6ef8-c18b-4f29-b413-e5055648c1cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caeff07d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow.keras.layers import Dense, Flatten, Conv2D\n",
    "\n",
    "\n",
    "import bentoml\n",
    "\n",
    "print(\"TensorFlow version:\", tf.__version__)\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train, x_test = x_train / 255.0, x_test / 255.0\n",
    "\n",
    "# Add a channels dimension\n",
    "x_train = x_train.reshape(60000, 28, 28, 1).astype(\"float32\")\n",
    "x_test = x_test.reshape(10000, 28, 28, 1).astype(\"float32\")\n",
    "\n",
    "train_ds = (\n",
    "    tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)\n",
    ")\n",
    "\n",
    "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)\n",
    "\n",
    "\n",
    "class MyModel(tf.Module):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.conv1 = Conv2D(32, 3, activation=\"relu\")\n",
    "        self.flatten = Flatten()\n",
    "        self.d1 = Dense(128, activation=\"relu\")\n",
    "        self.d2 = Dense(10)\n",
    "\n",
    "    @tf.function(input_signature=[tf.TensorSpec([None, 28, 28, 1], tf.float32)])\n",
    "    def __call__(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.d1(x)\n",
    "        return self.d2(x)\n",
    "\n",
    "\n",
    "# Create an instance of the model\n",
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38888f0a",
   "metadata": {},
   "source": [
    "## Training and Saving the model\n",
    "\n",
    "Then we initialize some simple tensorflow helper functions and create the training and testing methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c62db15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "train_loss = tf.keras.metrics.Mean(name=\"train_loss\")\n",
    "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name=\"train_accuracy\")\n",
    "\n",
    "test_loss = tf.keras.metrics.Mean(name=\"test_loss\")\n",
    "test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name=\"test_accuracy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "788c19a0",
   "metadata": {},
   "source": [
    "### Training and Testing TF Steps\n",
    "\n",
    "Now we assemble our TensorFlow2 training and testing steps. We use @tf.function as the new way (a part of TensorFlow2) to initialize a TensorFlow session.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b2fdd72",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_step(images, labels):\n",
    "    with tf.GradientTape() as tape:\n",
    "        # training=True is only needed if there are layers with different\n",
    "        # behavior during training versus inference (e.g. Dropout).\n",
    "        predictions = model(images)\n",
    "        loss = loss_object(labels, predictions)\n",
    "    gradients = tape.gradient(loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "\n",
    "    train_loss(loss)\n",
    "    train_accuracy(labels, predictions)\n",
    "\n",
    "\n",
    "def test_step(images, labels):\n",
    "    # training=False is only needed if there are layers with different\n",
    "    # behavior during training versus inference (e.g. Dropout).\n",
    "    predictions = model(images)\n",
    "    t_loss = loss_object(labels, predictions)\n",
    "\n",
    "    test_loss(t_loss)\n",
    "    test_accuracy(labels, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2104a6",
   "metadata": {},
   "source": [
    "### Training the model\n",
    "\n",
    "As provided by TensorFlow, we train and test the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d311c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 2\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    # Reset the metrics at the start of the next epoch\n",
    "    train_loss.reset_states()\n",
    "    train_accuracy.reset_states()\n",
    "    test_loss.reset_states()\n",
    "    test_accuracy.reset_states()\n",
    "\n",
    "    for images, labels in train_ds:\n",
    "        train_step(images, labels)\n",
    "\n",
    "    for test_images, test_labels in test_ds:\n",
    "        test_step(test_images, test_labels)\n",
    "\n",
    "    print(\n",
    "        f\"Epoch {epoch + 1}, \"\n",
    "        f\"Loss: {train_loss.result()}, \"\n",
    "        f\"Accuracy: {train_accuracy.result() * 100}, \"\n",
    "        f\"Test Loss: {test_loss.result()}, \"\n",
    "        f\"Test Accuracy: {test_accuracy.result() * 100}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92d9b23c",
   "metadata": {},
   "source": [
    "### Saving the model\n",
    "\n",
    "Finally, we make one call to the bentoml library to save this tensorflow model to be used later as part of the prediction service that we will create."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe9c4a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "bentoml.tensorflow.save_model(\n",
    "    \"tensorflow_mnist\",\n",
    "    model,\n",
    "    signatures={\"__call__\": {\"batchable\": True, \"batch_dim\": 0}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf35e55",
   "metadata": {},
   "source": [
    "## Create a BentoML Service for serving the model\n",
    "\n",
    "Note: using `%%writefile` here because `bentoml.Service` instance must be created in a separate `.py` file\n",
    "\n",
    "Even though we have only one model, we can create as many api endpoints as we want. Here we create two end points `predict_ndarray` and `predict_image`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e2f590",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile service.py\n",
    "\n",
    "import bentoml\n",
    "import numpy as np\n",
    "from bentoml.io import Image, NumpyNdarray\n",
    "from PIL.Image import Image as PILImage\n",
    "\n",
    "mnist_runner = bentoml.tensorflow.get(\"tensorflow_mnist:latest\").to_runner()\n",
    "\n",
    "svc = bentoml.Service(\n",
    "    name=\"tensorflow_mnist_demo\",\n",
    "    runners=[mnist_runner],\n",
    ")\n",
    "\n",
    "@svc.api(input=Image(), output=NumpyNdarray(dtype=\"float32\"))\n",
    "async def predict_image(f: PILImage) -> \"np.ndarray\":\n",
    "    assert isinstance(f, PILImage)\n",
    "    arr = np.array(f)/255.0\n",
    "    assert arr.shape == (28, 28)\n",
    "\n",
    "    # We are using greyscale image and our PyTorch model expect one\n",
    "    # extra channel dimension\n",
    "    arr = np.expand_dims(arr, (0, 3)).astype(\"float32\") # reshape to [1, 28, 28, 1]\n",
    "    return await mnist_runner.async_run(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "590147aa",
   "metadata": {},
   "source": [
    "Start a dev model server to test out the service defined above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "29173871",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "^C\n"
     ]
    }
   ],
   "source": [
    "!bentoml serve service.py:svc --reload"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "606c1b36",
   "metadata": {},
   "source": [
    "Now you can use something like:\n",
    "\n",
    "`curl -H \"Content-Type: multipart/form-data\" -F'fileobj=@samples/0.png;type=image/png' http://127.0.0.1:3000/predict_image`\n",
    "    \n",
    "to send an image to the digit recognition service."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6372c692",
   "metadata": {},
   "source": [
    "We can also do a simple local benchmark if [locust](https://locust.io/) is installed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a88e9879",
   "metadata": {},
   "outputs": [],
   "source": [
    "!locust --headless -u 500 -r 10 --run-time 10m --host http://127.0.0.1:3000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7f03564",
   "metadata": {},
   "source": [
    "## Build a Bento for distribution and deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f51f7406",
   "metadata": {},
   "source": [
    "A `bentofile` is already created in this directory for building a Bento for the service:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3757297b",
   "metadata": {},
   "source": [
    "```yaml\n",
    "service: \"service:svc\"\n",
    "description: \"file: ./README.md\"\n",
    "labels:\n",
    "  owner: bentoml-team\n",
    "  stage: demo\n",
    "include:\n",
    "- \"*.py\"\n",
    "exclude:\n",
    "- \"tests/\"\n",
    "python:\n",
    "  lock_packages: False\n",
    "  packages:\n",
    "    - tensorflow\n",
    "    - Pillow\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa8e50a",
   "metadata": {},
   "source": [
    "Note that we exclude `tests/` from the bento using exclude.\n",
    "\n",
    "Simply run `bentoml build` from current directory to build a Bento with the latest version of the `tensorflow_mnist` model. This may take a while when running for the first time for BentoML to resolve all dependency versions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d37b339",
   "metadata": {},
   "outputs": [],
   "source": [
    "!bentoml build"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36306933",
   "metadata": {},
   "source": [
    "Starting a dev server with the Bento build:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec4b9dff",
   "metadata": {},
   "outputs": [],
   "source": [
    "!bentoml serve tensorflow2_demo:latest"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "kernelspec": {
   "display_name": "tf2-py3.7",
   "language": "python",
   "name": "tf2-py3.7"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "name": "tensorflow2_mnist_demo.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
